neededDocumentDict = {}

internalIDWithExternalIDDict = {}

# for me to extract the needed documents for the training set
inputFileName0 = "/data3/obukai/workspace/web-search-engine-wei/polyIRIndexer/small_set_of_gov2_documents_needed_to_be_extracted_for_training_queries_sorted_by_trecID"
inputFileHandler0 = open(inputFileName0,"r")
for line in inputFileHandler0.readlines():
    lineElements = line.strip().split(" ")
    if lineElements[4] not in neededDocumentDict:
        neededDocumentDict[ lineElements[4] ] = 1
    else:
        print "NOT good"
        exit(1)

print "len(neededDocumentDict)",len(neededDocumentDict)

inputFileName = "/data3/obukai/workspace/web-search-engine-wei/polyIRIndexer/tempProcessingTrainingFileDocuments"
inputFileHandler = open(inputFileName,"r")
line = inputFileHandler.readline()

while line:
    lineElements = line.strip().split(" ")
    
    if lineElements[0] == "doc_id(internal):":
        currentInternalDocID = lineElements[1]
        nextLine = inputFileHandler.readline()
        nextLineElements = nextLine.strip().split(" ")
        if nextLineElements[1] not in neededDocumentDict:
            pass
        else:
            print nextLineElements[1],currentInternalDocID
            if nextLineElements[1] not in internalIDWithExternalIDDict:
                internalIDWithExternalIDDict[ nextLineElements[1] ] = currentInternalDocID
            else:
                print "NOT as expected."
                exit(1)
    
    line = inputFileHandler.readline()

print "len(internalIDWithExternalIDDict):",len(internalIDWithExternalIDDict)

inputFileName1 = "/data3/obukai/workspace/web-search-engine-wei/polyIRIndexer/small_set_of_gov2_documents_needed_to_be_extracted_for_training_queries_sorted_by_trecID"
inputFileHandler1 = open(inputFileName1,"r")

outputFileName = "/data3/obukai/workspace/web-search-engine-wei/polyIRIndexer/small_set_of_gov2_documents_needed_to_be_extracted_for_training_queries_sorted_by_trecID_NEW"
outputFileHandler = open(outputFileName,"w")
for line in inputFileHandler1.readlines():
    lineElements = line.strip().split(" ")
    
    queryID = lineElements[0] 
    rankInResultList = int( lineElements[1] )
    overallScore = float( lineElements[2] )
    NAPart = lineElements[3]
    trecID = lineElements[4]
    newInternalDocID = internalIDWithExternalIDDict[trecID]
    
    outputFileHandler.write(queryID + " " + str(rankInResultList) + " " + str(overallScore) + " " + newInternalDocID + " " + trecID + "\n")
    
inputFileHandler0.close()
inputFileHandler.close()
inputFileHandler1.close()

outputFileHandler.close()

